package cn.ixiaoxin.deepchat.websocket;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

@Slf4j
@ServerEndpoint(value = "/websocket/{userId}")
@Component
public class WebSocketServer {
    private static final Map<String, Session> sessions = new ConcurrentHashMap<>();
    private static final Map<String, String> userStatus = new ConcurrentHashMap<>();

    private static final long HEARTBEAT_INTERVAL = 30000; // 30秒
    private ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        sessions.put(userId, session);
        userStatus.put(userId, "online");
        broadcastUserStatus(userId, "online");
        broadcastAllUserStatus(userId);

        // 开始心跳检测
        scheduler.scheduleAtFixedRate(() -> {
            try {
                if (session.isOpen()) {
                    session.getBasicRemote().sendText("{\"type\":\"heartbeat\"}");
                }
            } catch (IOException e) {
                System.out.println("Heartbeat failed for user " + userId);
                try {
                    session.close();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
        }, HEARTBEAT_INTERVAL, HEARTBEAT_INTERVAL, TimeUnit.MILLISECONDS);
    }

    @OnMessage
    public void onMessage(String message, Session session, @PathParam("userId") String userId) {
        JSONObject json = JSON.parseObject(message);

        switch (json.getString("type")) {
            case "single":
                handleSingleMessage(userId, json);
                break;
            case "call":
                handleCallRequest(userId, json);
                break;
            case "call_response":
                handleCallResponse(userId, json);
                break;
            case "heartbeat":
                break;
            case "call_offer":
                forwardToTarget(userId, json);
                break;
            case "call_answer":
                forwardToTarget(userId, json);
                break;
            case "call_ice_candidate":
                forwardToTarget(userId, json);
                break;
            case "call_end":
                forwardToTarget(userId, json);
                break;
            default:
                System.out.println("Unknown message type: " + json.getString("type"));
        }
    }

    public void handleSingleMessage(String fromId, JSONObject message) {
        String toId = message.getString("toId");
        Session toSession = sessions.get(toId);

        if (toSession != null && toSession.isOpen()) {
            try {
                message.put("fromId", fromId);
                toSession.getBasicRemote().sendText(message.toJSONString());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public void handleCallRequest(String fromId, JSONObject message) {
        String toId = message.getString("toId");
        Session toSession = sessions.get(toId);

        if (toSession != null && toSession.isOpen()) {
            try {
                message.put("fromId", fromId);
                toSession.getBasicRemote().sendText(message.toJSONString());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public void handleCallResponse(String fromId, JSONObject message) {
        String toId = message.getString("toId");
        Session toSession = sessions.get(toId);

        if (toSession != null && toSession.isOpen()) {
            try {
                message.put("fromId", fromId);
                toSession.getBasicRemote().sendText(message.toJSONString());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private void forwardToTarget(String fromId, JSONObject message) {
        String toId = message.getString("toId");
        Session toSession = sessions.get(toId);

        if (toSession != null && toSession.isOpen()) {
            try {
                message.put("fromId", fromId);
                toSession.getBasicRemote().sendText(message.toJSONString());
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @OnClose
    public void onClose(Session session,  @PathParam("userId") String userId) {
        sessions.remove(userId);
        userStatus.put(userId, "offline");
        broadcastUserStatus(userId, "offline");
        broadcastAllUserStatus(userId);

        // 停止心跳检测
        scheduler.shutdown();
    }

    private void broadcastUserStatus(String userId, String status) {
        JSONObject message = new JSONObject();
        message.put("type", "status");
        message.put("userId", userId);
        message.put("status", status);

        sessions.forEach((id, session) -> {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message.toJSONString());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    private void broadcastAllUserStatus(String userId) {
        JSONObject message = new JSONObject();
        message.put("type", "allUserStatus");
        message.put("users", userStatus);

        sessions.forEach((id, session) -> {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message.toJSONString());
                } catch (IOException e) {
                    e.printStackTrace();
               }
            }
        });
    }
}